'''
从图片文件夹和标注txt文件夹中构建yolo数据集(无data.yaml)
'''


import os
import shutil
import random

# 源目录
images_src = '/mnt/d/aiready/chonche/v2/enhance_dataset_generate/images'
labels_src = '/mnt/d/aiready/chonche/v2/enhance_dataset_generate/labels'

# 目标目录
dataset_dir = '/mnt/d/aiready/chonche/v2/dataset'
images_dst = os.path.join(dataset_dir, 'images')
labels_dst = os.path.join(dataset_dir, 'labels')

# 创建目标文件夹
for folder in [images_dst, labels_dst]:
    for subfolder in ['train', 'val']:
        os.makedirs(os.path.join(folder, subfolder), exist_ok=True)

# 获取所有图片文件名（不带扩展名）
image_files = [f for f in os.listdir(images_src) if f.endswith('.png')]
image_names = [os.path.splitext(f)[0] for f in image_files]

# 打乱顺序
random.shuffle(image_names)

# 计算验证集数量（20%）
val_count = int(len(image_names) * 0.2)


# 分成 val 和 train
val_names = image_names[:val_count]
train_names = image_names[val_count:]

# 定义拷贝函数
def copy_files(file_list, src_img, src_lbl, dst_img, dst_lbl, subfolder):
    for name in file_list:
        img_src_path = os.path.join(src_img, name + '.png')
        lbl_src_path = os.path.join(src_lbl, name + '.txt')
        
        img_dst_path = os.path.join(dst_img, subfolder, name + '.png')
        lbl_dst_path = os.path.join(dst_lbl, subfolder, name + '.txt')
        
        # 拷贝图片和标签
        shutil.copyfile(img_src_path, img_dst_path)
        shutil.copyfile(lbl_src_path, lbl_dst_path)

# 拷贝文件到对应目录
copy_files(val_names, images_src, labels_src, images_dst, labels_dst, 'val')
copy_files(train_names, images_src, labels_src, images_dst, labels_dst, 'train')

print('数据集划分完成！')
